import time
from tqdm import tqdm
from transformers import BertTokenizer, BertForSequenceClassification
from torch.utils.data import TensorDataset, DataLoader
import torch

# 加载预训练的 tokenizer 和模型
tokenizer = BertTokenizer.from_pretrained('model/chinese-macbert-base')
model = BertForSequenceClassification.from_pretrained('model/chinese-macbert-base', num_labels=3)

# 准备训练数据
texts = []
labels = []


def read_txt(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        # 读取文件内容
        content = file.read()
        return content


# 获取文件内容
file_content = read_txt('data/train.txt')
# 按照换行符拆分内容，并按照空格拆分每行
lines = file_content.split('\n')
for line in lines:
    if line:
        title, number = line.split("	")
        texts.append(title)
        labels.append(int(number))

# 定义优化器和损失函数
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
loss_fn = torch.nn.CrossEntropyLoss()

# 编码文本数据
encoded_inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
input_ids = encoded_inputs['input_ids']
attention_mask = encoded_inputs['attention_mask']
labels = torch.tensor(labels)

# 显存训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
input_ids = input_ids.to(device)
attention_mask = attention_mask.to(device)
labels = labels.to(device)

# 创建数据集和数据加载器
dataset = TensorDataset(input_ids, attention_mask, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 训练模型
model.train()
for epoch in range(5):
    total_loss = 0
    progress_bar = tqdm(enumerate(dataloader), total=len(dataloader), desc="训练进度")
    for step, batch in progress_bar:
        batch_input_ids = batch[0]
        batch_attention_mask = batch[1]
        batch_labels = batch[2]

        optimizer.zero_grad()
        outputs = model(batch_input_ids, attention_mask=batch_attention_mask, labels=batch_labels)
        # 计算误差
        loss = outputs.loss
        total_loss += loss.item()
        # 反向传播和优化
        loss.backward()
        optimizer.step()
        progress_bar.set_postfix({'误差': loss.item()})

    avg_loss = total_loss / len(dataloader)
    time.sleep(0.2)
    print(f'第{epoch + 1}次迭代, 平均误差: {avg_loss:.4f}')

# 保存模型
torch.save(model.state_dict(), 'train/trained_model.pt')
